﻿//-----------------------------------------------------------------------------------------------------------
// WBFSSync Project by Omega Frost 
// http://wbfssync.codeplex.com/
//
// WBFSSync is Licensed under the terms of the 
// Microsoft Reciprocal License (Ms-RL)
//-----------------------------------------------------------------------------------------------------------
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Runtime.InteropServices;
using WBFSe3.Native;
using WBFSe3.IO;

namespace WBFSe3.Streams
{
    //-------------------------------------------------------------------------------------------------------
    //
    //-------------------------------------------------------------------------------------------------------
    public class WbfsFileStream : WbfsStream
    {
        //---------------------------------------------------------------------------------------------------
        // Variables
        //---------------------------------------------------------------------------------------------------
        private IntPtr fileHandle;
        private String name;
        private FileAccess access;
        private long position;

        protected bool bSparse;
        protected bool bCanSparse;


        //---------------------------------------------------------------------------------------------------
        // Properties
        //---------------------------------------------------------------------------------------------------
        public virtual IntPtr FileHandle { get { return this.fileHandle; } }
        public override String Name { get { return this.name; } }

        public override bool CanRead { get { return (this.access | FileAccess.Read) != 0; } }
        public override bool CanWrite { get { return (this.access | FileAccess.Write) != 0; } }
        public override bool CanSeek { get { return true; } }

        public override long Length { get { return GetLength(); } }

        public override long Position
        {
            get { return this.position; }
            set { Seek(value, SeekOrigin.Begin); }
        }


        //---------------------------------------------------------------------------------------------------
        // Routines
        //---------------------------------------------------------------------------------------------------

        //---------------------------------------------------------------------------------------------------
        // Constructor
        //---------------------------------------------------------------------------------------------------
        public WbfsFileStream(string path, int sectorSize, FileMode mode)
            : base(sectorSize)
        {
            if (path == null)
                throw new ArgumentNullException();

            if (path.Length == 0)
                throw new ArgumentException();

            if (((int)mode < 1) || ((int)mode > 6))
                throw new ArgumentOutOfRangeException();

            bool exists = File.Exists(path);
            bool created = (mode == FileMode.Create) || (mode == FileMode.CreateNew) ||
                ((mode == FileMode.OpenOrCreate) && !exists);
            
            this.name = path;
            this.access = FileAccess.ReadWrite;
            this.fileHandle = WbfsNative.CreateFile(path, this.access,
            FileShare.None, IntPtr.Zero,mode, EFileAttributes.None, IntPtr.Zero);

            if (this.fileHandle == WbfsNative.InvalidHandle)
            {
                int error = Marshal.GetLastWin32Error();
                switch (error)
                {
                    default: throw new WbfsIOException(error);
                    case 2: throw new FileNotFoundException();
                    case 3: throw new DirectoryNotFoundException();
                    case 5: throw new UnauthorizedAccessException();
                    case 32: throw new SharingViolationException();
                }
            }

            GetSparseCaps();
            if (CanWrite && created)
                SetSparseState();
        }


        //---------------------------------------------------------------------------------------------------
        // Constructor
        //---------------------------------------------------------------------------------------------------
        public WbfsFileStream(string path, int sectorSize, FileMode mode, FileAccess access)
            : base(sectorSize)
        {
            if (path == null)
                throw new ArgumentNullException();

            if (path.Length == 0)
                throw new ArgumentException();

            if (((int)mode < 1) || ((int)mode > 6))
                throw new ArgumentOutOfRangeException();

            bool exists = File.Exists(path);
            bool created = (mode == FileMode.Create) || (mode == FileMode.CreateNew) ||
                ((mode == FileMode.OpenOrCreate) && !exists);

            this.name = path;
            this.access = access;
            this.fileHandle = WbfsNative.CreateFile(path, access,
                FileShare.None, IntPtr.Zero, mode, EFileAttributes.None, IntPtr.Zero);

            if (this.fileHandle == WbfsNative.InvalidHandle)
            {
                int error = Marshal.GetLastWin32Error();
                switch (error)
                {
                    default: throw new WbfsIOException(error);
                    case 2: throw new FileNotFoundException();
                    case 3: throw new DirectoryNotFoundException();
                    case 5: throw new UnauthorizedAccessException();
                    case 32: throw new SharingViolationException();
                }
            }

            GetSparseCaps();
            if (CanWrite && created)
                SetSparseState();
        }


        //---------------------------------------------------------------------------------------------------
        // Constructor
        //---------------------------------------------------------------------------------------------------
        public WbfsFileStream(string path, int sectorSize, FileMode mode, FileAccess access, FileShare share)
            : base(sectorSize)
        {
            if (path == null)
                throw new ArgumentNullException();

            if (path.Length == 0)
                throw new ArgumentException();

            if (((int)mode < 1) || ((int)mode > 6))
                throw new ArgumentOutOfRangeException();

            bool exists = File.Exists(path);
            bool created = (mode == FileMode.Create) || (mode == FileMode.CreateNew) ||
                ((mode == FileMode.OpenOrCreate) && !exists);

            this.name = path;
            this.access = access;
            this.fileHandle = WbfsNative.CreateFile(path, access, share,
                IntPtr.Zero, mode, EFileAttributes.None, IntPtr.Zero);

            if (this.fileHandle == WbfsNative.InvalidHandle)
            {
                int error = Marshal.GetLastWin32Error();
                switch (error)
                {
                    default: throw new WbfsIOException(error);
                    case 2: throw new FileNotFoundException();
                    case 3: throw new DirectoryNotFoundException();
                    case 5: throw new UnauthorizedAccessException();
                    case 32: throw new SharingViolationException();
                }
            }

            if (File.Exists(path))
            {
                GetSparseCaps();
                if (CanWrite && created)
                    SetSparseState();
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Get drive and sparse caps
        //---------------------------------------------------------------------------------------------------
        private void GetSparseCaps()
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            //Check the drive
            WbfsDriveInfo drive;
            WbfsDriveInfo.GetDrive(this.name, out drive);
            this.bCanSparse = drive.IsSparse;

            //Check the file
            this.bSparse = (File.GetAttributes(this.Name) &
                FileAttributes.SparseFile) != 0;
        }


        //---------------------------------------------------------------------------------------------------
        // Set sparse state
        //---------------------------------------------------------------------------------------------------
        public void SetSparseState()
        {
            if (this.bSparse || !this.bCanSparse)
                return;

            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            uint temp = 0;
            if (WbfsNative.DeviceIoControl(this.fileHandle, EIOControlCode.FsctlSetSparse,
                IntPtr.Zero, 0, IntPtr.Zero, 0, ref temp, IntPtr.Zero) == 0)
            {
                int error = Marshal.GetLastWin32Error();
                throw new WbfsIOException();
            }

            this.bSparse = true;
        }


        //---------------------------------------------------------------------------------------------------
        // Set a sparse region of the file
        //---------------------------------------------------------------------------------------------------
        public void SetSparseRegion(long from, long to)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            uint temp = 0;
            SZERODATAINF zdi = new SZERODATAINF();
            zdi.FileOffset = from;
            zdi.BeyondFinalZero = to;

            int szdi = Marshal.SizeOf(zdi);
            IntPtr pzdi = Marshal.AllocHGlobal(szdi);
            Marshal.StructureToPtr(zdi, pzdi, true);

            bool fail = (WbfsNative.DeviceIoControl(this.fileHandle, 
                EIOControlCode.FsctlSetZeroData, pzdi, (uint)szdi, 
                IntPtr.Zero, 0, ref temp, IntPtr.Zero) == 0);
            
            Marshal.FreeHGlobal(pzdi);
            if (fail) throw new WbfsIOException();
        }


        //---------------------------------------------------------------------------------------------------
        // Get the used regions
        //---------------------------------------------------------------------------------------------------
        public SRANGEBUFFER[] QueryAlloctedLayout()
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (!this.bSparse)
            {
                SRANGEBUFFER[] b = new SRANGEBUFFER[1];
                b[0].FileOffset = 0;
                b[0].Length = Length;
                return b;
            }

            // Start with 1024 spaces
            uint szrb = (uint)Marshal.SizeOf(typeof(SRANGEBUFFER));
            IntPtr prba = Marshal.AllocHGlobal((int)szrb * 1024);

            // Final array
            SRANGEBUFFER[] layout = new SRANGEBUFFER[0];

            // Query the whole file
            SRANGEBUFFER range;
            IntPtr prange;

            range.FileOffset = 0;
            range.Length = Length;
            prange = Marshal.AllocHGlobal((int)szrb);

            int r, n, i, start = 0;
            uint bytes = 0;

            do
            {
                // Send the range to the unmanaged memory
                Marshal.StructureToPtr(range, prange, true);

                // Expand the array
                Array.Resize(ref layout, layout.Length + 1024);

                // Retrieve the layout
                r = WbfsNative.DeviceIoControl(this.fileHandle, EIOControlCode.FsctlQueryAllocatedRanges,
                    prange, szrb, prba, 1024 * szrb, ref bytes, IntPtr.Zero);

                if (r != 0)
                {
                    n = (int)(bytes / szrb);

                    // Read the data from unmanaged memory
                    long p = prba.ToInt64();
                    for (i = start; i < start + n; i++)
                    {
                        layout[i] = (SRANGEBUFFER)Marshal.PtrToStructure(
                            (IntPtr)p, typeof(SRANGEBUFFER));
                        p += szrb;
                    }

                    start += n;
                    range.FileOffset = layout[start - 1].FileOffset + layout[start - 1].Length;
                    range.Length = Length - range.FileOffset;
                }

            } while (r == 234); //ERROR_MORE_DATA

            // Realocate the buffer
            Array.Resize(ref layout, start);

            // Free the unmanaged memory
            Marshal.FreeHGlobal(prba);
            Marshal.FreeHGlobal(prange);

            if (r == 0)
            {
                throw new WbfsIOException();
            }
            else
            {
                return layout;
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Get the sparse (real) size
        //---------------------------------------------------------------------------------------------------
        public long QuerySparseSize()
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (!bSparse) 
                return Length;

            // Initial size
            long size = 0;

            // Start with 1024 spaces
            uint szrb = (uint)Marshal.SizeOf(typeof(SRANGEBUFFER));
            IntPtr prba = Marshal.AllocHGlobal((int)szrb * 1024);

            // Query the whole file
            SRANGEBUFFER range;
            IntPtr prange;

            range.FileOffset = 0;
            range.Length = Length;
            prange = Marshal.AllocHGlobal((int)szrb);

            int r, n, i;
            SRANGEBUFFER buffer = new SRANGEBUFFER();
            uint bytes = 0;

            do
            {
                // Send the range to the unmanaged memory
                Marshal.StructureToPtr(range, prange, true);

                // Retrieve the layout
                r = WbfsNative.DeviceIoControl(this.fileHandle, EIOControlCode.FsctlQueryAllocatedRanges,
                    prange, szrb, prba, 1024 * szrb, ref bytes, IntPtr.Zero);

                if (r != 0)
                {
                    n = (int)(bytes / szrb);

                    // Read the data from unmanaged memory
                    long p = prba.ToInt64();
                    for (i = 0; i < n; i++)
                    {
                        buffer = (SRANGEBUFFER)Marshal.
                            PtrToStructure((IntPtr)p,
                            typeof(SRANGEBUFFER));

                        p += szrb;
                        size += buffer.Length;
                    }

                    range.FileOffset = buffer.FileOffset + buffer.Length;
                    range.Length = Length - range.FileOffset;
                }

            } while (r == 234); // ERROR_MORE_DATA

            // Free the unmanaged memory
            Marshal.FreeHGlobal(prba);
            Marshal.FreeHGlobal(prange);

            if (r == 0)
            {
                throw new WbfsIOException();
            }
            else
            {
                return size;
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Writes a block of zeroes or mark the region as sparse
        //---------------------------------------------------------------------------------------------------
        public override void WriteEmptyBlock(int count)
        {
            if (count < 0)
                throw new ArgumentException();

            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (this.bSparse)
            {
                long from = this.Position;
                long to = from + count;

                SetSparseRegion(from, to);
                Seek(to, SeekOrigin.Begin);
            }
            else
            {
                Byte[] zeroes = new Byte[count];
                Array.Clear(zeroes, 0, zeroes.Length);
                Write(zeroes, 0, count);
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Signals the stream closing
        //---------------------------------------------------------------------------------------------------
        public override void Close()
        {
            base.Close();
        }


        //---------------------------------------------------------------------------------------------------
        // Gets the length of the stream
        //---------------------------------------------------------------------------------------------------
        protected long GetLength()
        {
            long length;
            long dummy;

            if (!WbfsNative.SetFilePointerEx(this.fileHandle, 0, out length, SeekOrigin.End) ||
            !WbfsNative.SetFilePointerEx(this.fileHandle, this.position, out dummy, SeekOrigin.Begin))
                throw new WbfsIOException();

            return length;
        }


        //---------------------------------------------------------------------------------------------------
        // Sets the length of the stream
        //---------------------------------------------------------------------------------------------------
        public override void SetLength(long value)
        {
            if (value < 0)
                throw new ArgumentOutOfRangeException();

            if (!this.CanRead && !this.CanWrite)
                throw new NotSupportedException();

            if (!WbfsNative.SetEndOfFile(this.fileHandle))
                throw new WbfsIOException();
        }


        //---------------------------------------------------------------------------------------------------
        // Releases the file handle
        //---------------------------------------------------------------------------------------------------
        protected override void Dispose(bool disposing)
        {
            if (this.fileHandle != WbfsNative.InvalidHandle)
            {
                WbfsNative.CloseHandle(this.fileHandle);
                this.fileHandle = WbfsNative.InvalidHandle;
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Flushes the file buffers
        //---------------------------------------------------------------------------------------------------
        public override void Flush()
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (!WbfsNative.FlushFileBuffers(this.fileHandle))
                throw new WbfsIOException();
        }


        //---------------------------------------------------------------------------------------------------
        // Sets the file pointer to a specific position
        //---------------------------------------------------------------------------------------------------
        public override long Seek(long offset, SeekOrigin origin)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            switch(origin)
            {
                case SeekOrigin.Current: offset += this.position; break;
                case SeekOrigin.End: offset = this.Length - offset; break;
            }

            if (offset == this.position)
                return this.position;

            if(offset < 0)
                throw new ArgumentException();

            if (!WbfsNative.SetFilePointerEx(this.fileHandle, offset, 
                out this.position, SeekOrigin.Begin))
                throw new WbfsIOException();

            return offset;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads data from the file
        //---------------------------------------------------------------------------------------------------
        public override int Read(byte[] array, int offset, int count)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (array.Length < offset + count)
                throw new ArgumentException();

            if (array == null)
                throw new ArgumentNullException();

            if ((offset < 0) || (count < 0))
                throw new ArgumentOutOfRangeException();

            if (!this.CanRead)
                throw new NotSupportedException();

            uint read;
            Byte[] buffer;

            if (offset == 0) buffer = array;
            else buffer = new Byte[count];

            if (!WbfsNative.ReadFile(this.fileHandle, buffer, 
                (uint)count, out read, IntPtr.Zero))
                throw new WbfsIOException();

            this.position += read;

            if (offset != 0) 
                buffer.CopyTo(array, offset);

            return (int)read;
        }


        //---------------------------------------------------------------------------------------------------
        // Writes data to the file
        //---------------------------------------------------------------------------------------------------
        public override void Write(byte[] array, int offset, int count)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (array.Length < offset + count)
                throw new ArgumentException();

            if (array == null)
                throw new ArgumentNullException();

            if ((offset < 0) || (count < 0))
                throw new ArgumentOutOfRangeException();

            if (!this.CanWrite)
                throw new NotSupportedException();

            uint written;
            Byte[] buffer;

            if (offset == 0) 
                buffer = array;
            else
            {
                buffer = new Byte[count];
                Array.Copy(array, offset, 
                    buffer, 0, count);
            }

            if (!WbfsNative.WriteFile(this.fileHandle, buffer,
                    (uint)count, out written, IntPtr.Zero))
                throw new WbfsIOException();

            this.position += written;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads aligned data from the file
        //---------------------------------------------------------------------------------------------------
        public override int ReadSector(int sector, IntPtr array, int offset)
        {
            uint read;

            if (!WbfsNative.SetFilePointerEx(this.fileHandle,
                (long)this.sectorSize * sector,
                out this.position, SeekOrigin.Begin))
                throw new WbfsIOException();

            if (!WbfsNative.ReadFile(this.fileHandle, IntPtr.Add(array, offset),
                (uint)this.sectorSize, out read, IntPtr.Zero))
                throw new WbfsIOException();

            this.position += this.sectorSize;
            return (int)read;
        }


        //---------------------------------------------------------------------------------------------------
        // Writes aligned data to the file
        //---------------------------------------------------------------------------------------------------
        public override int WriteSector(int sector, IntPtr array, int offset)
        {
            uint written;

            if (!WbfsNative.SetFilePointerEx(this.fileHandle,
                (long)this.sectorSize * sector,
                out this.position, SeekOrigin.Begin))
                throw new WbfsIOException();

            if (!WbfsNative.WriteFile(this.fileHandle, IntPtr.Add(array, offset),
                    (uint)this.sectorSize, out written, IntPtr.Zero))
                throw new WbfsIOException();

            this.position += this.sectorSize;

            return (int)written;
        }
    }
}
